import os
import platform
from transformers import AutoTokenizer, AutoModel
# coding=utf-8
import socket
import sys
import time
import struct
import threading


# path1
path1 = r'chatglm-6b\models--THUDM--chatglm-6b\snapshots\d2bbc82a2cdd04522ad340bdd464379808676950'
path2 = r"chatglm-6b-int4\models--THUDM--chatglm-6b-int4\snapshots\3ba943724188bca4a34c7cfc3f2752c05cf75ceb"


# tokenizer = AutoTokenizer.from_pretrained(path1,cache_dir="./chatglm-6b", trust_remote_code=True)
# model = AutoModel.from_pretrained(path1,cache_dir="./chatglm-6b", trust_remote_code=True).half().quantize(4).cuda()

# # 按需修改，目前只支持 4/8 bit 量化
# model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().quantize(4).cuda()

# tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b-int4",cache_dir="./chatglm-6b-int4", trust_remote_code=True)
# model = AutoModel.from_pretrained("THUDM/chatglm-6b-int4",cache_dir="./chatglm-6b-int4", trust_remote_code=True).half().cuda()

# 本地模型
tokenizer = AutoTokenizer.from_pretrained(path2,  trust_remote_code=True)
model = AutoModel.from_pretrained(path2, trust_remote_code=True).half().cuda()


model = model.eval()

'''
报文格式 ：  
ID  RID  总长度  类型  实际数据
类型定义：｛0：系统消息，1:文字，2：图片｝
其中ID都是 8 字节， 数据长度 4字节  类型1字节 共 21字节 
'''


class Server():
    head_format = '>QQIB'  # struct 头部格式
    head_len = 21   # 共 21字节
    code = 'utf-16'  # 文字编码
    user = dict()   # 当前在线用户

    # 接收长度为length的字节，返回该字节流
    def recv_len(self, sc, length):
        data = bytes()
        while len(data) < length:
            data += sc.recv(length - len(data))
        return data

    # 接收一个完整意义的数据包，包括头部和数据部分
    def recv_data(self, sc: socket):
        head_data = self.recv_len(sc, self.head_len)
        head = struct.unpack(self.head_format, head_data)
        data = self.recv_len(sc, head[2])
        return head, data

    # 将字节流打包发送
    def send_data(self, ID, RID, t, sc: socket, data: bytes):
        length = len(data)
        head = struct.pack(self.head_format, ID, RID, length, t)
        sc.sendall(head + data)
        # sc.sendall()

    def gen_resp(self, query:str,local_hist:list, sc:socket, isRuning:list ):
        def build_prompt(history ): # 直接发送最新的回答
            if not history: return ""
            query, response = history[-1] 
            return response   
        
        count = 0
        for response, tmp_his in model.stream_chat(tokenizer, query, history=local_hist):
            if not isRuning[0]:
                break 
            count += 1
            if count % 4 == 0:
                res_txt = build_prompt( tmp_his) 
                print(res_txt, end="\n", flush=True)
                self.send_data(0x896551, 0x888752, 1, sc, (res_txt).encode(self.code)) # 每次发送完整字符
            # 让 历史截取,防止输入溢出
            if tmp_his:
                local_hist.clear() 
                MaxLen = 0
                tmp_his.reverse()
                for q,v in tmp_his:
                    MaxLen += len(q) + len(v)
                    # print("MaxLen", MaxLen)
                    if MaxLen < 1836:
                        local_hist.append( (q,v) )  # 保证引用  
                    else:
                        break
                local_hist.reverse() # 保证引用  
            
        # os.system(clear_command)
        res_txt = build_prompt(local_hist )
        print(res_txt, end=" ]]\n")
        self.send_data(0x612d1231, 0x48462, 15, sc, (res_txt+"[END]").encode(self.code)) # type=15表示结束
        isRuning[0] = False


    # 处理sc发来的数据，每次处理一个数据包
    def handle(self, sc: socket, local_hist:list):
        isRuning = [False] 
        while (True):
            try:
                head, data = self.recv_data(sc)
                #  head = [ ID  RID  总长度  类型 ]
                # print("收到一个消息：")
                # print(head)
                cli_txt = data.decode(self.code)
                # 处理的代码
                print("\n\n用户：", cli_txt)
                query = cli_txt

                if query == "break":
                    isRuning[0] = False
                    continue
                if not isRuning[0]:

                    if query == "stop":
                        return False, 0
                    if query == "clear":
                        local_hist.clear()
                        continue
                     
                    isRuning[0] = True
                    t = threading.Thread(target=self.gen_resp, args=(query, local_hist,sc, isRuning))
                    t.setDaemon(True)
                    t.start()   
                # 处理的代码
            
            except Exception as e:
                print(e)
                return False, 0
        return True, 0

    def _handle_all(self, sc):
        local_hist = [] 
        self.handle(sc, local_hist)  

    # 开启服务器
    def server(self, ip, port):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.bind((ip, port))
        sock.listen(5)
        print('开始监听：', sock.getsockname())
        while True:
            sc, addr = sock.accept()
            del addr
            th = threading.Thread(target=self._handle_all, args=(sc,))
            th.setDaemon(True)
            th.start()


def task():
    Server().server('0.0.0.0', 1060)


if __name__ == "__main__":
    t = threading.Thread(target=task)
    t.daemon = True
    t.start()

    while True:
        try:
            time.sleep(1)
        except KeyboardInterrupt:
            print('收到键盘信号 Ctrl-C，Bye！')
            exit(0)
